from osgeo import gdal,osr,gdalconst,ogr
from xml.dom.minidom import parse
import datetime as dt
import json
import daemon
import redis
import psycopg2
import shutil
import zipfile
from rocketmq.client import Producer, Message
import os
import re


def name_2_info(cargs,name):
    connection = psycopg2.connect(**cargs)
    connection.autocommit = True
    cursor = connection.cursor()
    query = f"""select tile,product_date,cloud,tif_path,ST_AsText(geom)
        from sentinel2_l1c
        where name = '{name}';
        """
    cursor.execute(query)
    data = cursor.fetchone()
    cursor.close()
    connection.close()
    return data


def read_xmlfile(xml_path,select_node):
    domTree = parse(xml_path)
    collection = domTree.documentElement
    select_element = [collection]
    for node_dict in select_node:
        tag_name = node_dict["tag_name"]
        if node_dict["attribute"] is None:
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
            ]
        else:
            attribute_name = node_dict["attribute"]["name"]
            attribute_value = node_dict["attribute"]["value"]
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
                if element_sub.getAttribute(attribute_name) == attribute_value
            ]
    datas = [
        element.firstChild.data
        for element in select_element
    ]
    return datas


def read_cloud(zip_path):
    select_node = [
        {
            "tag_name": "n1:Quality_Indicators_Info",
            "attribute": None
        },
        {
            "tag_name": "Cloud_Coverage_Assessment",
            "attribute": None
        }
    ]
    xml_namere = "^S2.*\.SAFE/MTD_MSIL1C\.xml$"
    zip_file = zipfile.ZipFile(zip_path, "r")
    cloud = "00"
    for file_info in zip_file.infolist():
        if re.match(xml_namere, file_info.filename):
            xml_path = zip_file.open(file_info.filename, "r")
            datas = read_xmlfile(xml_path, select_node)
            cloud = round(float(datas[0]))
            cloud = f"0{cloud}" if len(f"{cloud}") == 1 else f"{cloud}"
    return float(cloud)


# 发送数据到消息队列
def send_data(json_dir,zip_path,tif_path,redisr,topic="henan_rsanalysis_s2_1c_collect",addr='192.168.1.70:9876'):
    name = os.path.basename(zip_path).replace(".zip","")
    tile =  name.split("_")[5][1:]
    product_date = name.split("_")[2][:8].replace("-","")
    # producer = Producer("data")
    # producer.set_name_server_address(addr)
    # producer.start()
    data = {
        "tile": tile,  # 假设其他属性字段在结果中的索引为1
        "product_date": product_date,
        "cpjbdm": "1",
        "cpjdmc": "L1C",
        "imgtype_code": "S2",   
        "imgtype_name": "Sentinel-2",
        "cloud": float(read_cloud(zip_path)),
        "tif_path": tif_path,
        "geometry": tif_2_geo(tif_path)}
    json_path = os.path.join(json_dir,f"{name}.json")
    with open(json_path,"w") as json_file:
        json.dump(data,json_file,indent=4)
    redisr.rpush('raster_json_upload', f'"{json_path}"')
    # message_content = {"file_path":json_path}
    # message_body = json.dumps(message_content)
    # msg = Message(topic)
    # msg.set_tags("TagA")
    # msg.set_body(message_body)
    # _ = producer.send_sync(msg)
    # producer.shutdown()
    return 


def zip_2_jp2(zip_path, bands, save_dir):
    band_names = {
        "1":"T.*_B01\.jp2",
        "2":"T.*_B02\.jp2",
        "3":"T.*_B03\.jp2",
        "4":"T.*_B04\.jp2",
        "5":"T.*_B05\.jp2",
        "6":"T.*_B06\.jp2",
        "7":"T.*_B07\.jp2",
        "8A":"T.*_B8A\.jp2",
        "8":"T.*_B08\.jp2",
        "9":"T.*_B09\.jp2",
        "10":"T.*_B10\.jp2",
        "11":"T.*_B11\.jp2",
        "12":"T.*_B12\.jp2"}
    pathres = list(map(lambda x:".*/IMG_DATA/"+band_names[f"{x}"],bands))
    paths = []
    zipFile = zipfile.ZipFile(zip_path, "r")
    for pathre in pathres:
        for info in zipFile.infolist():
            if re.match(pathre,info.filename):
                path = zipFile.extract(info.filename, save_dir)
                paths.append(path)
    zipFile.close()
    return paths
    

def jp2s_2_tif(jp2_paths,tif_path,nodata=0):
    vrt_path = tif_path.replace('.tif','.vrt')
    options = gdal.BuildVRTOptions(
        resolution="highest",
        VRTNodata=nodata,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        separate=True)
    gdal.BuildVRT(vrt_path,jp2_paths,options=options)
    creationOptions = [
        "BIGTIFF=YES",
        "TILED=YES",
        "BLOCKXSIZE=1024",
        # "TFW=YES",
        "NUM_THREADS=ALL_CPUS"]
    options = gdal.TranslateOptions(
        format="GTiff",
        xRes=10,
        yRes=10,
        # dstNodata=0,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        creationOptions=creationOptions)
    gdal.Translate(tif_path,vrt_path,options=options)
    os.remove(vrt_path)
    return


def zip_2_tif(zip_path,save_dir,bands:list=['2','3','4','5','6','7',"8A",'8','11','12']):
    zip_name = os.path.basename(zip_path)
    year_month = zip_name.split("_")[2][:6]
    date_dir = os.path.join(save_dir,year_month)
    if not os.path.exists(date_dir):os.makedirs(date_dir)
    tif_name = zip_name.replace('.SAFE', '').replace('.zip', '.tif')
    tif_path = os.path.join(date_dir, tif_name)
    if os.path.exists(tif_path): return tif_path
    jp2_paths = zip_2_jp2(zip_path,bands,date_dir)
    jp2s_2_tif(jp2_paths,tif_path)
    unzip_dir = tif_path.replace(".tif",".SAFE")
    if os.path.exists(unzip_dir): shutil.rmtree(unzip_dir)
    return tif_path


def projection_2_spatil(tif_path):
    # 打开栅格数据集
    dataset = gdal.Open(tif_path)
    # 获取WKT格式的空间参考信息
    wkt = dataset.GetProjection()
    # 创建空间参考对象
    spatial_ref = osr.SpatialReference(wkt=wkt)
    # 检查空间参考是否存在
    if wkt == '':
        raise "无空间参考信息"
    else:
        # 检查是否为地理坐标系
        if spatial_ref.IsGeographic():
            print("使用的是地理坐标系")
            # 获取EPSG代码
            epsg_code = spatial_ref.GetAttrValue('AUTHORITY', 1)
            # 检查是否是WGS84
            if epsg_code == '4326':
                print("该栅格使用的是WGS84坐标系。")
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
            else:
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                # 使用gdal.Warp()进行重新投影
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
        # 检查是否为投影坐标系
        elif spatial_ref.IsProjected():
            print("使用的是投影坐标系")
            # 使用gdal.Warp()进行重新投影
            # vrt_path = re.sub("\..*$",".vrt",tif_path)
            mem_path = '/vsimem/output.tif'
            gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
            return mem_path
        
        
def tif_2_geo(tif_path):
    # 打开栅格数据集
    raster_path = projection_2_spatil(tif_path)
    dataset = gdal.Open(raster_path,gdalconst.GA_Update)
    wkt_projection = dataset.GetProjection()
    band = dataset.GetRasterBand(1)
    nodata = band.GetNoDataValue()
    array = band.ReadAsArray()
    mask_band = band.GetMaskBand()
    mask_array = mask_band.ReadAsArray()
    band.WriteArray(mask_array)
    band = dataset.GetRasterBand(1)
    srs = osr.SpatialReference()
    srs.ImportFromWkt(wkt_projection)
    # 创建输出矢量数据源
    driver = ogr.GetDriverByName('Memory') # Memory
    mem_ds = driver.CreateDataSource('memory') # memory
    layer = mem_ds.CreateLayer('polygon_layer', srs, ogr.wkbPolygon)
    # 添加字段
    fieldDef = ogr.FieldDefn('bandvalue', ogr.OFTInteger)
    layer.CreateField(fieldDef)
    options = []
    options.append('8CONNECTED=8')  # 使用 8 连通区域生长算法
    gdal.Polygonize(
        band, 
        mask_band, 
        layer, 
        -1, 
        options=options, 
        callback=gdal.TermProgress_nocb
    )
    fearture = layer.GetFeature(0) 
    geometry = fearture.geometry()
    geometry_wkt = geometry.ExportToWkt()
    mem_ds.Destroy()
    del dataset
    return geometry_wkt


def main():
    with open('/data/fengyy/project/CopernicusSentinel2Level1C/main/config.json',"r") as file:
        params = json.load(file)
    json_dir = params['json_dir']
    tif_dir = params["tif_dir"]
    rediscon = params["rediscon"]
    pool = redis.ConnectionPool(**rediscon)
    redisr = redis.Redis(connection_pool=pool)
    while True:
        _,element = redisr.blpop('raster_geometry',0)
        zip_path = element.decode('utf-8').replace('"','')
        print(zip_path)
        tif_path = zip_2_tif(zip_path, tif_dir)
        send_data(json_dir,zip_path,tif_path,redisr)


def run_daemon():
    with daemon.DaemonContext(
        stdout=open(r"/data/Yang/data/logfile/RasterGeometry.log","w"),
        stderr=open(r"/data/Yang/data/logfile/RasterGeometry.log","w")
        ):
        main()
    return


if __name__ == "__main__":
    # main()
    zip_path = r"/data/ygjsb/data/ProvinceMonthZip/41HeNan/202401/S2B_MSIL1C_20240104T030119_N0510_R032_T50SKC_20240104T044342.zip"
    # zip_path = r"/data/ygjsb/data/ProvinceMonthZip/41HeNan/202404/S2A_MSIL1C_20240421T030531_N0510_R075_T50SKC_20240421T064424.zip"
    tif_dir = r"/data/fengyy/dateset/temp"
    zip_2_tif(zip_path, tif_dir)
    # run_daemon()